# 创建一个简单的"去不去看电影"决策模型
# 影响因素: [天气(0=下雨,1=晴天), 有空闲时间(小时), 票价(元), 想看程度(0-10)]

import torch

'''

# 准备训练数据
X = torch.tensor([
    [0, 2, 50, 8],   # 下雨天,2小时空闲,票价50,想看程度8 => 去看(1)
    [1, 3, 40, 7],   # 晴天,3小时空闲,票价40,想看程度7 => 去看(1)
    [0, 1, 60, 3],   # 下雨天,1小时空闲,票价60,想看程度3 => 不去看(0)
    [1, 2, 45, 5],   # 晴天,2小时空闲,票价45,想看程度5 => 去看(1)
    [0, 3, 35, 2],   # 下雨天,3小时空闲,票价35,想看程度2 => 不去看(0)
], dtype=torch.float32)

y = torch.tensor([[1], [1], [0], [1], [0]], dtype=torch.float32)

# 创建一个简单的神经网络模型
model = torch.nn.Sequential(
    torch.nn.Linear(4, 8),
    torch.nn.ReLU(),
    torch.nn.Linear(8, 1),
    torch.nn.Sigmoid()
)

# 定义损失函数和优化器
criterion = torch.nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(1000):
    # 前向传播
    y_pred = model(X)
    loss = criterion(y_pred, y)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch % 100 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

# 测试模型
test_case = torch.tensor([1, 2, 45, 9], dtype=torch.float32)  # 晴天,2小时空闲,票价45,想看程度9
prediction = model(test_case)
print(f"\n是否去看电影? (概率: {prediction.item():.2f})")
if prediction.item() > 0.5:
    print("建议: 去看!")
else:
    print("建议: 不去看!")

'''

class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = torch.nn.Linear(4, 8)
        self.relu = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(8, 1)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

def train_model(model, X, y, epochs=1000, lr=0.01):
    criterion = torch.nn.BCELoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    for epoch in range(epochs):
        optimizer.zero_grad()
        y_pred = model(X)
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()

        if epoch % 100 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item():.4f}')

# 准备训练数据
X = torch.tensor([
    [0, 2, 50, 8],   # 下雨天,2小时空闲,票价50,想看程度8 => 去看(1)
    [1, 3, 40, 7],   # 晴天,3小时空闲,票价40,想看程度7 => 去看(1)
    [0, 1, 60, 3],   # 下雨天,1小时空闲,票价60,想看程度3 => 不去看(0)
    [1, 2, 45, 5],   # 晴天,2小时空闲,票价45,想看程度5 => 去看(1)
    [0, 3, 35, 2],   # 下雨天,3小时空闲,票价35,想看程度2 => 不去看(0)
], dtype=torch.float32)

y = torch.tensor([[1], [1], [0], [1], [0]], dtype=torch.float32)

model = MyModel()
train_model(model, X, y, epochs=1000, lr=0.01)

# 测试模型
test_case = torch.tensor([1, 2, 45, 9], dtype=torch.float32)  # 晴天,2小时空闲,票价45,想看程度9
prediction = model(test_case)
print(f"\n是否去看电影? (概率: {prediction.item():.2f})")
if prediction.item() > 0.5:
    print("建议: 去看!")
else:
    print("建议: 不去看!")
